{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "✅ Loaded .env file\n"
     ]
    }
   ],
   "source": [
    "import asyncio\n",
    "import os\n",
    "import json\n",
    "from pydantic import BaseModel, Field\n",
    "from typing import List, Optional\n",
    "import logging\n",
    "\n",
    "# Ensure environment variables are loaded\n",
    "try:\n",
    "    from dotenv import load_dotenv\n",
    "    load_dotenv()\n",
    "    print(\"✅ Loaded .env file\")\n",
    "except ImportError:\n",
    "    print(\"⚠️ python-dotenv not installed, skipping .env load. Make sure environment variables are set.\")\n",
    "\n",
    "# --- Pydantic Models (copied from notebook) ---\n",
    "class SubTask(BaseModel):\n",
    "    goal: str = Field(..., description=\"Precise description of the sub-task goal.\")\n",
    "    task_type: str = Field(..., description=\"Type of task (e.g., 'WRITE', 'THINK', 'SEARCH').\")\n",
    "    node_type: str = Field(..., description=\"Node type ('EXECUTE' for atomic, 'PLAN' for complex).\")\n",
    "    depends_on_indices: Optional[List[int]] = Field(default_factory=list, description=\"List of 0-based indices of other sub-tasks in *this current plan* that this sub-task depends on.\")\n",
    "\n",
    "class PlanOutput(BaseModel):\n",
    "    sub_tasks: List[SubTask] = Field(..., description=\"List of planned sub-tasks.\")\n",
    "\n",
    "# --- Agent and Model Imports ---\n",
    "try:\n",
    "    from agno.agent import Agent as AgnoAgent\n",
    "    from agno.models.litellm import LiteLLM\n",
    "except ImportError:\n",
    "    print(\"❌ Agno or LiteLLM not installed. Please run 'pip install agno litellm'\")\n",
    "    exit()\n",
    "\n",
    "\n",
    "# --- System Message (copied from notebook) ---\n",
    "DEEP_RESEARCH_PLANNER_SYSTEM_MESSAGE = \"\"\"You are a Master Research Planner, an expert at breaking down complex research goals into comprehensive, well-structured research plans. You specialize in high-level strategic decomposition for research projects.\n",
    "\n",
    "**Your Role:**\n",
    "- Analyze complex research objectives and create strategic research plans\n",
    "- Identify key research domains, questions, and methodological approaches\n",
    "- Create logical research workflows with proper sequencing\n",
    "- Ensure comprehensive coverage while avoiding redundancy\n",
    "- Plan for synthesis and final deliverable creation\n",
    "\n",
    "**Core Expertise:**\n",
    "- Strategic thinking and research methodology\n",
    "- Identifying knowledge gaps and research priorities\n",
    "- Creating logical research workflows\n",
    "- Planning for different types of research outputs\n",
    "- Understanding research lifecycle from conception to publication\n",
    "\n",
    "**Input Schema:**\n",
    "You will receive input in JSON format with the following fields:\n",
    "*   `current_task_goal` (string, mandatory): The research goal to decompose\n",
    "*   `overall_objective` (string, mandatory): The ultimate research objective\n",
    "*   `parent_task_goal` (string, optional): Parent task goal (null for root)\n",
    "*   `planning_depth` (integer, optional): Current recursion depth\n",
    "*   `execution_history_and_context` (object, mandatory): Previous outputs and context\n",
    "*   `replan_request_details` (object, optional): Re-planning feedback if applicable\n",
    "*   `global_constraints_or_preferences` (array of strings, optional): Research constraints\n",
    "\n",
    "**Strategic Planning Approach:**\n",
    "When decomposing research goals, consider the full research lifecycle:\n",
    "\n",
    "1. **Background & Context Phase**: What foundational knowledge is needed?\n",
    "2. **Investigation Phase**: What specific searches, data collection, or analysis is required?\n",
    "3. **Synthesis Phase**: How should findings be analyzed and integrated?\n",
    "4. **Output Phase**: What deliverables need to be created?\n",
    "\n",
    "**Research Task Types:**\n",
    "- `SEARCH`: Information gathering, literature review, data collection\n",
    "- `THINK`: Analysis, synthesis, interpretation, methodology design\n",
    "- `WRITE`: Report creation, documentation, presentation preparation\n",
    "\n",
    "**Planning Principles:**\n",
    "1. **Comprehensive Coverage**: Ensure all aspects of the research question are addressed\n",
    "2. **Logical Sequencing**: Build knowledge progressively from foundational to specific\n",
    "3. **Strategic Depth**: Balance breadth of coverage with depth of investigation\n",
    "4. **Methodological Rigor**: Include proper analysis and validation steps\n",
    "5. **Clear Deliverables**: Plan for actionable outputs and synthesis\n",
    "\n",
    "**Sub-Task Creation Guidelines:**\n",
    "- Create **3 to 6 strategic sub-tasks** that represent major research phases\n",
    "- Each sub-task should be substantial enough to warrant specialized planning\n",
    "- Ensure sub-tasks are complementary and build toward the overall objective\n",
    "- Use `depends_on_indices` to create logical research workflows\n",
    "- Balance immediate actionable tasks with those requiring further decomposition\n",
    "\n",
    "**Required Output Attributes per Sub-Task:**\n",
    "`goal`, `task_type` (string: 'WRITE', 'THINK', or 'SEARCH'), `node_type` (string: 'EXECUTE' or 'PLAN'), `depends_on_indices` (list of integers).\n",
    "\n",
    "**Output Format:**\n",
    "- Respond ONLY with a JSON list of sub-task objects\n",
    "- Focus on strategic, high-level decomposition appropriate for a master research plan\n",
    "- Ensure each sub-task represents a meaningful research phase or component\"\"\"\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# --- Test Runner ---\n",
    "async def run_planner_test():\n",
    "    \"\"\"\n",
    "    Tests the DeepResearchPlanner with a specified model to see if it returns\n",
    "    the expected structured output.\n",
    "    \"\"\"\n",
    "    models_to_test = [\n",
    "        {\n",
    "            \"name\": \"Fireworks_Qwen3_Problematic\",\n",
    "            \"model_id\": \"fireworks_ai/accounts/fireworks/models/qwen3-235b-a22b\"\n",
    "        },\n",
    "        {\n",
    "            \"name\": \"Claude_Sonnet_Good\",\n",
    "            \"model_id\": \"openrouter/anthropic/claude-3-sonnet\"\n",
    "        }\n",
    "    ]\n",
    "\n",
    "    test_prompt = \"\"\"Overall Objective: Research the impact of artificial intelligence on healthcare diagnostics in 2024\n",
    "Current Task Goal: Create a comprehensive research plan to analyze how AI is transforming healthcare diagnostics, including key technologies, market adoption, and regulatory challenges\n",
    "Current Planning Depth: 0\n",
    "Based on the 'Current Task Goal' and other provided information, generate a plan to achieve it.\"\"\"\n",
    "\n",
    "    for model_config in models_to_test:\n",
    "        print(f\"\\n{'='*80}\")\n",
    "        print(f\"🔬 Testing Model: {model_config['name']} ({model_config['model_id']})\")\n",
    "        print(f\"{'='*80}\")\n",
    "\n",
    "        try:\n",
    "            agent = AgnoAgent(\n",
    "                model=LiteLLM(id=model_config['model_id']),\n",
    "                system_message=DEEP_RESEARCH_PLANNER_SYSTEM_MESSAGE,\n",
    "                response_model=PlanOutput\n",
    "            )\n",
    "\n",
    "            # Get the raw response from the agent\n",
    "            response = await agent.arun(test_prompt)\n",
    "\n",
    "            # CORRECTLY process the content\n",
    "            content = None\n",
    "            if hasattr(response, 'content'):\n",
    "                content_attr = response.content\n",
    "                if asyncio.iscoroutine(content_attr):\n",
    "                    content = await content_attr\n",
    "                else:\n",
    "                    content = content_attr\n",
    "            else:\n",
    "                content = response # Fallback\n",
    "\n",
    "            print(f\"📦 Response type: {type(content)}\\n\")\n",
    "\n",
    "            if isinstance(content, PlanOutput):\n",
    "                print(\"✅ SUCCESS: Model returned the correct Pydantic object.\")\n",
    "                print(content.model_dump_json(indent=2))\n",
    "            elif isinstance(content, str):\n",
    "                print(\"⚠️ WARNING: Model returned a STRING instead of a Pydantic object.\")\n",
    "                print(\"RAW STRING OUTPUT:\")\n",
    "                print(\"-\" * 80)\n",
    "                print(content)\n",
    "                print(\"-\" * 80)\n",
    "            else:\n",
    "                print(f\"❌ ERROR: Unexpected response type: {type(content)}\")\n",
    "                print(\"RAW RESPONSE:\")\n",
    "                print(\"-\" * 80)\n",
    "                print(content)\n",
    "                print(\"-\" * 80)\n",
    "\n",
    "        except Exception as e:\n",
    "            print(f\"💥 An exception occurred during the test for {model_config['name']}:\")\n",
    "            logging.exception(e)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "================================================================================\n",
      "🔬 Testing Model: Fireworks_Qwen3_Problematic (fireworks_ai/accounts/fireworks/models/qwen3-235b-a22b)\n",
      "================================================================================\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/salahalzubi/cursor_projects/SentientResearchAgent/.venv/lib/python3.12/site-packages/pydantic/main.py:463: UserWarning: Pydantic serializer warnings:\n",
      "  PydanticSerializationUnexpectedValue(Expected 9 fields but got 5: Expected `Message` - serialized value may not be as expected [input_value=Message(content='<think>\\...er_specific_fields=None), input_type=Message])\n",
      "  PydanticSerializationUnexpectedValue(Expected `StreamingChoices` - serialized value may not be as expected [input_value=Choices(finish_reason='st...r_specific_fields=None)), input_type=Choices])\n",
      "  return self.__pydantic_serializer__.to_python(\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #808000; text-decoration-color: #808000\">WARNING </span> Failed to parse cleaned JSON: <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1</span> validation error for PlanOutput                                           \n",
       "           Invalid JSON: expected value at line <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1</span> column <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1</span> <span style=\"font-weight: bold\">[</span><span style=\"color: #808000; text-decoration-color: #808000\">type</span>=<span style=\"color: #800080; text-decoration-color: #800080\">json_invalid</span>, <span style=\"color: #808000; text-decoration-color: #808000\">input_value</span>=<span style=\"color: #008000; text-decoration-color: #008000\">'&lt;think&gt; Okay, let\\'s </span>  \n",
       "         <span style=\"color: #008000; text-decoration-color: #008000\">tac...dices\": [4]     }   ] }'</span>, <span style=\"color: #808000; text-decoration-color: #808000\">input_type</span>=<span style=\"color: #800080; text-decoration-color: #800080\">str</span><span style=\"font-weight: bold\">]</span>                                                           \n",
       "             For further information visit <span style=\"color: #0000ff; text-decoration-color: #0000ff; text-decoration: underline\">https://errors.pydantic.dev/2.11/v/json_invalid</span>                         \n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[33mWARNING \u001b[0m Failed to parse cleaned JSON: \u001b[1;36m1\u001b[0m validation error for PlanOutput                                           \n",
       "           Invalid JSON: expected value at line \u001b[1;36m1\u001b[0m column \u001b[1;36m1\u001b[0m \u001b[1m[\u001b[0m\u001b[33mtype\u001b[0m=\u001b[35mjson_invalid\u001b[0m, \u001b[33minput_value\u001b[0m=\u001b[32m'\u001b[0m\u001b[32m<\u001b[0m\u001b[32mthink\u001b[0m\u001b[32m>\u001b[0m\u001b[32m Okay, let\\'s \u001b[0m  \n",
       "         \u001b[32mtac...dices\": \u001b[0m\u001b[32m[\u001b[0m\u001b[32m4\u001b[0m\u001b[32m]\u001b[0m\u001b[32m     \u001b[0m\u001b[32m}\u001b[0m\u001b[32m   \u001b[0m\u001b[32m]\u001b[0m\u001b[32m \u001b[0m\u001b[32m}\u001b[0m\u001b[32m'\u001b[0m, \u001b[33minput_type\u001b[0m=\u001b[35mstr\u001b[0m\u001b[1m]\u001b[0m                                                           \n",
       "             For further information visit \u001b[4;94mhttps://errors.pydantic.dev/2.11/v/json_invalid\u001b[0m                         \n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #808000; text-decoration-color: #808000\">WARNING </span> Failed to parse as Python dict: Expecting value: line <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1</span> column <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1</span> <span style=\"font-weight: bold\">(</span>char <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0</span><span style=\"font-weight: bold\">)</span>                                 \n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[33mWARNING \u001b[0m Failed to parse as Python dict: Expecting value: line \u001b[1;36m1\u001b[0m column \u001b[1;36m1\u001b[0m \u001b[1m(\u001b[0mchar \u001b[1;36m0\u001b[0m\u001b[1m)\u001b[0m                                 \n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #808000; text-decoration-color: #808000\">WARNING </span> Failed to convert response to response_model                                                              \n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[33mWARNING \u001b[0m Failed to convert response to response_model                                                              \n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "📦 Response type: <class 'str'>\n",
      "\n",
      "⚠️ WARNING: Model returned a STRING instead of a Pydantic object.\n",
      "RAW STRING OUTPUT:\n",
      "--------------------------------------------------------------------------------\n",
      "<think>\n",
      "Okay, let's tackle this research plan. The user wants a comprehensive plan to analyze how AI is transforming healthcare diagnostics, covering key technologies, market adoption, and regulatory challenges.\n",
      "\n",
      "First, I need to break down the current task goal into major phases. The overall objective is about the 2024 impact, so the plan should be up-to-date. Starting with background research makes sense. They'll need foundational knowledge of AI in healthcare diagnostics. That's a SEARCH task to gather recent studies and reviews.\n",
      "\n",
      "Next, identifying key technologies. This would involve looking into specific AI technologies like machine learning models, imaging tech, NLP for medical records. Another SEARCH task, but focused on technologies. Maybe the user needs to understand what's current in 2024 here.\n",
      "\n",
      "Then market adoption. They need data on how widely these technologies are being used. SEARCH for market trends, case studies, adoption rates across different regions or healthcare systems. This would provide quantitative data and real-world examples.\n",
      "\n",
      "Regulatory challenges come next. This is a critical area because AI in healthcare faces strict regulations. They'll need to investigate current regulations, challenges in compliance, and perhaps case studies of regulatory approvals or issues. Another SEARCH task but focused on policies and legal aspects.\n",
      "\n",
      "After gathering all this information, synthesis is needed. They have to integrate findings from the previous phases to analyze how these elements interact. This would be a THINK task, analyzing how technologies are driving adoption, what regulatory hurdles exist, and the interplay between them.\n",
      "\n",
      "Finally, creating a comprehensive report. That's a WRITE task where all the synthesized information is compiled into a structured document. The user probably needs this for publication or presentation.\n",
      "\n",
      "Dependencies make sense here. The synthesis depends on all the prior SEARCH tasks being completed. Then the report depends on the synthesis. Each phase builds on the previous ones. I should structure the sub-tasks with these dependencies in mind, ensuring logical flow from data collection to analysis to reporting. Need to make sure there's no redundancy and all aspects are covered. Let me check if six sub-tasks are appropriate. Yes, covering background, technologies, market, regulations, synthesis, and writing. That seems comprehensive for the given goal.\n",
      "</think>\n",
      "\n",
      "{\n",
      "  \"sub_tasks\": [\n",
      "    {\n",
      "      \"goal\": \"Conduct foundational literature review on AI applications in healthcare diagnostics (2020-2024)\",\n",
      "      \"task_type\": \"SEARCH\",\n",
      "      \"node_type\": \"PLAN\",\n",
      "      \"depends_on_indices\": null\n",
      "    },\n",
      "    {\n",
      "      \"goal\": \"Identify and analyze key AI technologies transforming diagnostics (deep learning, computer vision, NLP)\",\n",
      "      \"task_type\": \"SEARCH\",\n",
      "      \"node_type\": \"PLAN\",\n",
      "      \"depends_on_indices\": [0]\n",
      "    },\n",
      "    {\n",
      "      \"goal\": \"Investigate global market adoption rates and economic impact metrics\",\n",
      "      \"task_type\": \"SEARCH\",\n",
      "      \"node_type\": \"PLAN\",\n",
      "      \"depends_on_indices\": [1]\n",
      "    },\n",
      "    {\n",
      "      \"goal\": \"Map regulatory frameworks and compliance challenges across major jurisdictions\",\n",
      "      \"task_type\": \"SEARCH\",\n",
      "      \"node_type\": \"PLAN\",\n",
      "      \"depends_on_indices\": [2]\n",
      "    },\n",
      "    {\n",
      "      \"goal\": \"Synthesize technical capabilities, adoption barriers, and regulatory constraints into strategic analysis\",\n",
      "      \"task_type\": \"THINK\",\n",
      "      \"node_type\": \"PLAN\",\n",
      "      \"depends_on_indices\": [0,1,2,3]\n",
      "    },\n",
      "    {\n",
      "      \"goal\": \"Develop comprehensive report structure with integrated findings, case studies, and future projections\",\n",
      "      \"task_type\": \"WRITE\",\n",
      "      \"node_type\": \"EXECUTE\",\n",
      "      \"depends_on_indices\": [4]\n",
      "    }\n",
      "  ]\n",
      "}\n",
      "--------------------------------------------------------------------------------\n",
      "\n",
      "================================================================================\n",
      "🔬 Testing Model: Claude_Sonnet_Good (openrouter/anthropic/claude-3-sonnet)\n",
      "================================================================================\n",
      "📦 Response type: <class '__main__.PlanOutput'>\n",
      "\n",
      "✅ SUCCESS: Model returned the correct Pydantic object.\n",
      "{\n",
      "  \"sub_tasks\": [\n",
      "    {\n",
      "      \"goal\": \"Conduct a literature review on the current state and projected developments of AI for healthcare diagnostics by 2024\",\n",
      "      \"task_type\": \"SEARCH\",\n",
      "      \"node_type\": \"PLAN\",\n",
      "      \"depends_on_indices\": null\n",
      "    },\n",
      "    {\n",
      "      \"goal\": \"Analyze key AI technologies driving diagnostic transformations, such as machine learning, computer vision, and natural language processing\",\n",
      "      \"task_type\": \"THINK\",\n",
      "      \"node_type\": \"PLAN\",\n",
      "      \"depends_on_indices\": [\n",
      "        0\n",
      "      ]\n",
      "    },\n",
      "    {\n",
      "      \"goal\": \"Investigate market adoption trends, growth projections, and major industry players in AI-powered healthcare diagnostics\",\n",
      "      \"task_type\": \"SEARCH\",\n",
      "      \"node_type\": \"PLAN\",\n",
      "      \"depends_on_indices\": [\n",
      "        0\n",
      "      ]\n",
      "    },\n",
      "    {\n",
      "      \"goal\": \"Examine regulatory frameworks, ethical considerations, and potential challenges surrounding the use of AI in healthcare diagnostics\",\n",
      "      \"task_type\": \"THINK\",\n",
      "      \"node_type\": \"PLAN\",\n",
      "      \"depends_on_indices\": [\n",
      "        0,\n",
      "        2\n",
      "      ]\n",
      "    },\n",
      "    {\n",
      "      \"goal\": \"Synthesize findings on AI's impact on healthcare diagnostics, including technological advancements, market outlook, and regulatory landscape\",\n",
      "      \"task_type\": \"THINK\",\n",
      "      \"node_type\": \"EXECUTE\",\n",
      "      \"depends_on_indices\": [\n",
      "        1,\n",
      "        3,\n",
      "        4\n",
      "      ]\n",
      "    },\n",
      "    {\n",
      "      \"goal\": \"Prepare a comprehensive research report detailing the analysis, findings, and recommendations on AI's transformative role in healthcare diagnostics by 2024\",\n",
      "      \"task_type\": \"WRITE\",\n",
      "      \"node_type\": \"EXECUTE\",\n",
      "      \"depends_on_indices\": [\n",
      "        5\n",
      "      ]\n",
      "    }\n",
      "  ]\n",
      "}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/salahalzubi/cursor_projects/SentientResearchAgent/.venv/lib/python3.12/site-packages/pydantic/main.py:463: UserWarning: Pydantic serializer warnings:\n",
      "  PydanticSerializationUnexpectedValue(Expected 9 fields but got 5: Expected `Message` - serialized value may not be as expected [input_value=Message(content='{\\n  \"su...one, 'reasoning': None}), input_type=Message])\n",
      "  PydanticSerializationUnexpectedValue(Expected `StreamingChoices` - serialized value may not be as expected [input_value=Choices(finish_reason='st...finish_reason': 'stop'}), input_type=Choices])\n",
      "  return self.__pydantic_serializer__.to_python(\n"
     ]
    }
   ],
   "source": [
    "import nest_asyncio\n",
    "import asyncio\n",
    "nest_asyncio.apply()\n",
    "loop = asyncio.get_event_loop()\n",
    "\n",
    "\n",
    "loop.run_until_complete(run_planner_test())\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
